"""
Middleware components for CodeMCP Gateway.

Implements CORS support, rate limiting, authentication hooks,
and other request/response processing middleware.
"""

import time
import hashlib
from typing import Dict, Any, Optional, List, Callable
from collections import defaultdict, deque
from datetime import datetime, timedelta

from fastapi import FastAPI, Request, Response, HTTPException, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.trustedhost import TrustedHostMiddleware
from fastapi.middleware.gzip import GZipMiddleware
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from starlette.middleware.base import BaseHTTPMiddleware

from ..core.config import Config
from ..core.error_handler import log_info, log_warning, log_debug, handle_error


class RateLimitMiddleware(BaseHTTPMiddleware):
    """Rate limiting middleware using token bucket algorithm."""
    
    def __init__(self, 
                 app,
                 calls_per_minute: int = 60,
                 burst_size: int = 10,
                 enable_ip_based: bool = True,
                 enable_user_based: bool = True):
        super().__init__(app)
        self.calls_per_minute = calls_per_minute
        self.burst_size = burst_size
        self.enable_ip_based = enable_ip_based
        self.enable_user_based = enable_user_based
        
        # Rate limit storage: key -> (tokens, last_refill)
        self.buckets: Dict[str, tuple] = {}
        
        # Cleanup old entries periodically
        self.last_cleanup = time.time()
        self.cleanup_interval = 300  # 5 minutes
    
    async def dispatch(self, request: Request, call_next):
        """Process request with rate limiting."""
        # Skip rate limiting for health checks and documentation
        if request.url.path in ["/health", "/docs", "/redoc", "/openapi.json"]:
            return await call_next(request)
        
        # Determine rate limit key
        rate_limit_key = self._get_rate_limit_key(request)
        
        # Check rate limit
        if not self._check_rate_limit(rate_limit_key):
            log_warning(f"Rate limit exceeded for key: {rate_limit_key[:20]}...")
            raise HTTPException(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                detail={
                    "error": "Rate limit exceeded",
                    "limit": f"{self.calls_per_minute} calls per minute",
                    "retry_after": 60
                }
            )
        
        # Process request
        response = await call_next(request)
        
        # Add rate limit headers
        remaining = self._get_remaining_tokens(rate_limit_key)
        response.headers["X-RateLimit-Limit"] = str(self.calls_per_minute)
        response.headers["X-RateLimit-Remaining"] = str(max(0, remaining))
        response.headers["X-RateLimit-Reset"] = str(int(time.time()) + 60)
        
        return response
    
    def _get_rate_limit_key(self, request: Request) -> str:
        """Generate rate limit key from request."""
        # Try to get user ID from session or auth header
        user_id = None
        
        # Check for session header
        session_id = request.headers.get("X-Session-ID")
        if session_id:
            user_id = f"session:{session_id}"
        
        # Check for authorization header
        if not user_id:
            auth_header = request.headers.get("Authorization")
            if auth_header:
                # Create hash of auth header for privacy
                user_id = f"auth:{hashlib.sha256(auth_header.encode()).hexdigest()[:16]}"
        
        # Fall back to IP address
        if not user_id and self.enable_ip_based:
            client_ip = self._get_client_ip(request)
            user_id = f"ip:{client_ip}"
        
        return user_id or "anonymous"
    
    def _get_client_ip(self, request: Request) -> str:
        """Extract client IP address from request."""
        # Check for forwarded headers (reverse proxy)
        forwarded_for = request.headers.get("X-Forwarded-For")
        if forwarded_for:
            return forwarded_for.split(",")[0].strip()
        
        real_ip = request.headers.get("X-Real-IP")
        if real_ip:
            return real_ip
        
        # Fallback to direct client IP
        if request.client:
            return request.client.host
        
        return "unknown"
    
    def _check_rate_limit(self, key: str) -> bool:
        """Check if request is within rate limit."""
        current_time = time.time()
        
        # Cleanup old entries periodically
        if current_time - self.last_cleanup > self.cleanup_interval:
            self._cleanup_old_entries(current_time)
        
        # Get or create bucket
        if key not in self.buckets:
            self.buckets[key] = (self.burst_size, current_time)
            return True
        
        tokens, last_refill = self.buckets[key]
        
        # Calculate tokens to add based on time elapsed
        time_elapsed = current_time - last_refill
        tokens_to_add = (time_elapsed / 60) * self.calls_per_minute
        
        # Update token count (max = burst_size)
        new_tokens = min(self.burst_size, tokens + tokens_to_add)
        
        if new_tokens >= 1:
            # Allow request and consume token
            self.buckets[key] = (new_tokens - 1, current_time)
            return True
        else:
            # Rate limit exceeded
            self.buckets[key] = (new_tokens, current_time)
            return False
    
    def _get_remaining_tokens(self, key: str) -> int:
        """Get remaining tokens for a key."""
        if key not in self.buckets:
            return self.burst_size
        
        tokens, _ = self.buckets[key]
        return int(tokens)
    
    def _cleanup_old_entries(self, current_time: float):
        """Remove old rate limit entries."""
        cutoff_time = current_time - 3600  # Remove entries older than 1 hour
        
        keys_to_remove = []
        for key, (_, last_refill) in self.buckets.items():
            if last_refill < cutoff_time:
                keys_to_remove.append(key)
        
        for key in keys_to_remove:
            del self.buckets[key]
        
        self.last_cleanup = current_time
        
        if keys_to_remove:
            log_debug(f"Cleaned up {len(keys_to_remove)} old rate limit entries")


class RequestLoggingMiddleware(BaseHTTPMiddleware):
    """Middleware for detailed request/response logging."""
    
    def __init__(self, app, enable_detailed_logging: bool = False):
        super().__init__(app)
        self.enable_detailed_logging = enable_detailed_logging
    
    async def dispatch(self, request: Request, call_next):
        """Log request and response details."""
        start_time = time.time()
        
        # Log request
        if self.enable_detailed_logging:
            log_debug(f"Request: {request.method} {request.url}")
            log_debug(f"Headers: {dict(request.headers)}")
        
        # Process request
        try:
            response = await call_next(request)
            
            # Calculate processing time
            processing_time = time.time() - start_time
            
            # Log response
            log_info(f"{request.method} {request.url.path} - "
                    f"Status: {response.status_code} - "
                    f"Time: {processing_time:.3f}s")
            
            # Add timing header
            response.headers["X-Processing-Time"] = f"{processing_time:.3f}"
            
            return response
            
        except Exception as e:
            processing_time = time.time() - start_time
            
            log_warning(f"{request.method} {request.url.path} - "
                       f"Error: {str(e)} - "
                       f"Time: {processing_time:.3f}s")
            raise


class SecurityHeadersMiddleware(BaseHTTPMiddleware):
    """Middleware to add security headers to responses."""
    
    def __init__(self, app):
        super().__init__(app)
        self.security_headers = {
            "X-Content-Type-Options": "nosniff",
            "X-Frame-Options": "DENY",
            "X-XSS-Protection": "1; mode=block",
            "Referrer-Policy": "strict-origin-when-cross-origin",
            "Content-Security-Policy": "default-src 'self'",
        }
    
    async def dispatch(self, request: Request, call_next):
        """Add security headers to response."""
        response = await call_next(request)
        
        # Add security headers
        for header, value in self.security_headers.items():
            response.headers[header] = value
        
        return response


class AuthenticationMiddleware(BaseHTTPMiddleware):
    """Authentication middleware with hooks for custom auth providers."""
    
    def __init__(self, 
                 app, 
                 enabled: bool = False,
                 auth_provider: Optional[Callable] = None,
                 exclude_paths: List[str] = None):
        super().__init__(app)
        self.enabled = enabled
        self.auth_provider = auth_provider
        self.exclude_paths = exclude_paths or [
            "/", "/health", "/docs", "/redoc", "/openapi.json"
        ]
        self.bearer_scheme = HTTPBearer(auto_error=False)
    
    async def dispatch(self, request: Request, call_next):
        """Authenticate request if authentication is enabled."""
        if not self.enabled:
            return await call_next(request)
        
        # Skip authentication for excluded paths
        if request.url.path in self.exclude_paths:
            return await call_next(request)
        
        # Extract credentials
        credentials = await self.bearer_scheme(request)
        
        if not credentials:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Missing authentication credentials",
                headers={"WWW-Authenticate": "Bearer"},
            )
        
        # Validate credentials
        user_info = None
        if self.auth_provider:
            try:
                user_info = await self.auth_provider(credentials.credentials)
            except Exception as e:
                log_warning(f"Authentication failed: {e}")
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail="Invalid authentication credentials",
                    headers={"WWW-Authenticate": "Bearer"},
                )
        
        # Add user info to request state
        if user_info:
            request.state.user = user_info
        
        return await call_next(request)


def setup_middleware(app: FastAPI, config: Config):
    """
    Setup all middleware for the FastAPI application.
    
    Args:
        app: FastAPI application instance
        config: Configuration object
    """
    # Get middleware settings from config
    custom_settings = config.custom_settings
    
    # CORS Middleware
    cors_enabled = custom_settings.get("cors_enabled", True)
    if cors_enabled:
        allowed_origins = custom_settings.get("cors_allowed_origins", ["*"])
        allowed_methods = custom_settings.get("cors_allowed_methods", ["*"])
        allowed_headers = custom_settings.get("cors_allowed_headers", ["*"])
        
        app.add_middleware(
            CORSMiddleware,
            allow_origins=allowed_origins,
            allow_credentials=True,
            allow_methods=allowed_methods,
            allow_headers=allowed_headers,
        )
        log_info(f"CORS middleware enabled - Origins: {allowed_origins}")
    
    # Trusted Host Middleware (security)
    trusted_hosts = custom_settings.get("trusted_hosts", ["*"])
    if trusted_hosts and trusted_hosts != ["*"]:
        app.add_middleware(TrustedHostMiddleware, allowed_hosts=trusted_hosts)
        log_info(f"Trusted host middleware enabled - Hosts: {trusted_hosts}")
    
    # GZip Middleware (compression)
    gzip_enabled = custom_settings.get("gzip_enabled", True)
    if gzip_enabled:
        min_size = custom_settings.get("gzip_min_size", 1000)
        app.add_middleware(GZipMiddleware, minimum_size=min_size)
        log_info("GZip compression middleware enabled")
    
    # Security Headers Middleware
    security_headers_enabled = custom_settings.get("security_headers_enabled", True)
    if security_headers_enabled:
        app.add_middleware(SecurityHeadersMiddleware)
        log_info("Security headers middleware enabled")
    
    # Rate Limiting Middleware
    rate_limit_enabled = custom_settings.get("rate_limit_enabled", True)
    if rate_limit_enabled:
        calls_per_minute = custom_settings.get("rate_limit_calls_per_minute", 60)
        burst_size = custom_settings.get("rate_limit_burst_size", 10)
        
        app.add_middleware(
            RateLimitMiddleware,
            calls_per_minute=calls_per_minute,
            burst_size=burst_size
        )
        log_info(f"Rate limiting middleware enabled - "
                f"{calls_per_minute} calls/min, burst: {burst_size}")
    
    # Request Logging Middleware
    request_logging_enabled = custom_settings.get("request_logging_enabled", True)
    if request_logging_enabled:
        detailed_logging = config.debug or custom_settings.get("detailed_request_logging", False)
        app.add_middleware(RequestLoggingMiddleware, enable_detailed_logging=detailed_logging)
        log_info("Request logging middleware enabled")
    
    # Authentication Middleware
    auth_enabled = custom_settings.get("auth_enabled", False)
    if auth_enabled:
        # In a production system, you would implement a real auth provider
        # For now, we'll create a placeholder
        async def placeholder_auth_provider(token: str) -> Dict[str, Any]:
            """Placeholder authentication provider."""
            # This would validate the token against your auth system
            # For development, accept any token that starts with "dev-"
            if token.startswith("dev-"):
                return {"user_id": token, "role": "developer"}
            else:
                raise ValueError("Invalid token")
        
        exclude_paths = custom_settings.get("auth_exclude_paths", [
            "/", "/health", "/docs", "/redoc", "/openapi.json"
        ])
        
        app.add_middleware(
            AuthenticationMiddleware,
            enabled=auth_enabled,
            auth_provider=placeholder_auth_provider,
            exclude_paths=exclude_paths
        )
        log_info("Authentication middleware enabled")
    
    log_info("All middleware components configured successfully")


# Utility functions for use in route handlers

def get_user_info(request: Request) -> Optional[Dict[str, Any]]:
    """Get authenticated user information from request."""
    return getattr(request.state, "user", None)


def get_session_id(request: Request) -> Optional[str]:
    """Extract session ID from request headers."""
    return request.headers.get("X-Session-ID")


def get_client_info(request: Request) -> Dict[str, Any]:
    """Extract client information from request."""
    return {
        "ip": request.headers.get("X-Forwarded-For", 
                                  request.headers.get("X-Real-IP", 
                                                     getattr(request.client, "host", "unknown") if request.client else "unknown")),
        "user_agent": request.headers.get("User-Agent", "unknown"),
        "referer": request.headers.get("Referer"),
        "accept_language": request.headers.get("Accept-Language")
    }